import os
import random
import librosa
import soundfile as sf  # librosa推荐的音频写入库
from pathlib import Path


def cut_wav(input_path, output_path, min_length=8, max_length=20):
    """
    使用librosa处理音频，通过随机开始时间和长度裁剪WAV文件
    :param input_path: 输入WAV文件路径
    :param output_path: 输出文件夹路径
    :param min_length: 最小裁剪长度（秒）
    :param max_length: 最大裁剪长度（秒）
    """
    input_path = Path(input_path)
    print(f"处理文件: {input_path.name}")

    # sr=None 保留原始采样率，mono=True 转为单声道（如需保留立体声可设为False）
    y, sr = librosa.load(input_path, sr=None, mono=True)

    duration = librosa.get_duration(y=y, sr=sr)
    duration = int(duration)  # 转为整数秒

    # 确保最大裁剪长度不超过音频总时长
    max_possible_length = min(max_length, duration)

    # 随机生成裁剪长度（整数秒）
    clip_length = random.randint(min_length, max_possible_length)

    # 计算最大起始时间（确保裁剪不超出音频范围）
    max_start = duration - clip_length
    start_seconds = random.randint(0, max_start)  # 随机起始时间（整数秒）
    end_seconds = start_seconds + clip_length

    # 索引 = 时间（秒） * 采样率
    start_idx = int(start_seconds * sr)
    end_idx = int(end_seconds * sr)

    clipped_y = y[start_idx:end_idx]

    # 确保输出文件夹存在
    os.makedirs(output_path, exist_ok=True)
    # 生成输出文件名（包含原始文件名、起始时间和长度）
    file_name = input_path.stem
    output_file = Path(output_path) / f"{file_name} - {start_seconds}s - {clip_length}s.wav"

    # 使用soundfile写入WAV文件（保持原始采样率和格式）
    sf.write(
        file=output_file,
        data=clipped_y,
        samplerate=sr,
        subtype='PCM_16'  # 16位PCM编码（WAV常用格式）
    )

    print(f"裁剪完成：{output_file}（时长：{clip_length:.2f}秒）")


def cut_wav_s(folder_path, output_path, min_length=8, max_length=20):
    """
    批量处理文件夹中的WAV文件，使用随机开始时间和长度裁剪
    :param folder_path: 输入文件夹路径
    :param output_path: 输出文件夹路径
    :param min_length: 最小裁剪长度（秒）
    :param max_length: 最大裁剪长度（秒）
    """
    folder_path = Path(folder_path)
    # 遍历文件夹中所有WAV文件
    for wav_file in folder_path.glob("*.wav"):
        if wav_file.is_file():  # 确保是文件（非文件夹）
            cut_wav(
                input_path=str(wav_file),
                output_path=output_path,
                min_length=min_length,
                max_length=max_length
            )


if __name__ == "__main__":
    # input_wav = "data/稻香.wav"  # 单个文件路径
    input_folder = "../data/music_library"      # 批量处理的文件夹
    output_folder = "../data/test_data"   # 输出文件夹

    # 确保输出文件夹存在
    os.makedirs(output_folder, exist_ok=True)

    # 处理单个文件
    # cut_wav(input_wav, output_folder, min_length=10, max_length=20)

    # 批量处理文件夹中的文件
    cut_wav_s(input_folder, output_folder, min_length=8, max_length=20)
